import numpy as  np 
from skimage import io
from skimage import transform
import matplotlib.pyplot as plt 
import os
import torch
import torchvision
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import transforms
from torchvision.utils import make_grid
from PIL import Image
# label_map = {"on": 0, "off": 1}
# label_map = {"back_an": 0, "back_liang": 1, "front_an": 2, "front_liang": 3}
# label_map = {"CarForward": 0, "CarBackward": 1, "CarLeftward": 2, "CarRightward": 3, "OtherCarDirection": 4}
label_map = {'Extinguished': 0, 'RightRed': 1, 'StraightRed': 2, 'LeftRed': 3, 'RightYellow': 4, 'StraightYellow': 5,
             'LeftYellow': 6, 'RightGreen': 7, 'StraightGreen': 8, 'LeftGreen': 9, 'NonTrafficlight': 10, 'NonMoter': 11,
             'StraightRightRed': 12, 'StraightLeftRed': 13, 'StraightRightYellow': 14, 'StraightLeftYellow': 15,
             'StraightRightGreen': 16, 'StraightLeftGreen': 17, 'ForkLightRed': 18}
#熄灭的 # 右转红灯# 直行红灯# 左转红灯# 右转黄灯# 直行黄灯# 左转黄灯# 右转绿灯# 直行绿灯# 左转绿灯# 非信号灯# 非机动车信号灯
# 直行右转红灯# 直行左转红灯# 直行右转黄灯# 直行左转黄灯# 直行右转绿灯# 直行左转绿灯# 叉形灯}
class Resize(object):
    def __init__(self, output_size):
        self.output_size = output_size

    def __call__(self, sample):
        image = sample["image"]
        image_new = transform.resize(image, self.output_size)
        return {"image":image_new, "label":sample["label"]}

class ToTensor(object):

    def __call__(self, sample):
        image = sample["image"]
        image_new = np.transpose(image, (2,0,1))
        return {"image": torch.from_numpy(image_new), "label":sample["label"]}


class MyDataSet(Dataset):
    def __init__(self, root_dir,  transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.size = 0
        self.image_list = []
        self.label_list = []

        if not os.path.isdir(self.root_dir):
            print(self.root_dir + " dost not etist! ")

        for root, dirs, files in os.walk(root_dir):
            for file_name in files:
                file_root_split = os.path.join(root,file_name).strip().split('/')
                self.image_list.append(os.path.join(root,file_name))
                self.label_list.append(label_map.get(file_root_split[len(file_root_split)-2]))
                self.size += 1

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        image_path = self.image_list[idx]
        if not os.path.isfile(image_path):
            print(image_path + " does not exist! ")
            return None
        # image = io.imread(image_path)
        image = Image.open(image_path)
        label = int(self.label_list[idx])
        # sample = {"image":image, "label":label}
        # if self.transform:
        #     sample = self.transform(sample)
        if self.transform:
            image = self.transform(image)
        return image, label, image_path

if __name__ == "__main__":
    # train_dataset = MyDataSet(root_dir = '/media/data_1/Ubuntu/pytorch/cls_net/test2/data/train',\
    #     transform = transforms.Compose([Resize((64,64)), ToTensor()]))
    # trainset_dataloader = DataLoader(dataset = train_dataset, batch_size=4, shuffle=True, num_workers=4)
    train_dataset = MyDataSet(root_dir = '/media/pc/data_1/luo/chejiahao/pytorch/cls_net/data/wenben/test',\
        transform = transforms.Compose([
            # transforms.RandomCrop((64,64)),
            # transforms.RandomHorizontalFlip(),
            # transforms.RandomRotation((-5.0,5.0)),
            # transforms.Resize([48,112]),
            transforms.ToTensor(),
            # transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))
            transforms.Normalize((0.150, 0.126, 0.123), (0.245, 0.209, 0.201))
        ]))
    trainset_dataloader = DataLoader(dataset = train_dataset, batch_size=1, shuffle=True, num_workers=4)

    def show_images_batch(sample_batched):
        # images_batch, labels_batch = \
        # sample_batched['image'], sample_batched['label']
        images_batch, labels_batch, image_path= sample_batched
        print(images_batch.size())
        print(images_batch[0][0][0][0])
        print(images_batch[0][1][0][0])
        print(images_batch[0][2][0][0])
        grid = make_grid(images_batch)
        plt.imshow(grid.numpy().transpose(1, 2, 0))


    # sample_batch:  Tensor , NxCxHxW
    plt.figure()
    for i_batch, sample_batch in enumerate(trainset_dataloader):
        show_images_batch(sample_batch)
        plt.axis('off')
        plt.ioff()
        plt.show()


